package com.itao.ssh.server;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.itao.ssh.constant.ConstantPool;
import com.itao.ssh.pojo.SSHConnectInfo;
import com.itao.ssh.pojo.WebSSHData;
import com.jcraft.jsch.Channel;
import com.jcraft.jsch.ChannelShell;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.OnClose;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * @author sjt
 * @Date 2025/3/20 8:58
 */
@ServerEndpoint("/webssh")
@Component
@Slf4j
public class WebSocketServer {

    public static final Map<String, SSHConnectInfo> clients = new ConcurrentHashMap<>();
    //线程池
    private final ExecutorService executorService = Executors.newCachedThreadPool();

    @OnOpen
    public void onOpen(Session session) {
        JSch jSch = new JSch();
        SSHConnectInfo sshConnectInfo = new SSHConnectInfo();
        sshConnectInfo.setJSch(jSch);
        sshConnectInfo.setSession(session);
        String id = session.getId();
        //将这个ssh连接信息放入map中
        clients.put(id, sshConnectInfo);
    }

    @OnMessage
    public void onMessage(String buffer, Session session) {
        log.info("客户端发送数据:{}", buffer);
        ObjectMapper objectMapper = new ObjectMapper();
        WebSSHData webSSHData = null;
        try {
            webSSHData = objectMapper.readValue(buffer, WebSSHData.class);
        } catch (IOException e) {
            log.error("Json转换异常");
            log.error("异常信息:{}", e.getMessage());
            return;
        }
        String id = session.getId();
        if (ConstantPool.WEBSSH_OPERATE_CONNECT.equals(webSSHData.getOperate())) {
            //找到刚才存储的ssh连接对象
            SSHConnectInfo sshConnectInfo = clients.get(id);
            //启动线程异步处理
            WebSSHData finalWebSSHData = webSSHData;
            executorService.execute(() -> {
                try {
                    connectToSSH(sshConnectInfo, finalWebSSHData, session);
                    if (sshConnectInfo.getChannel().isClosed()) {
                        close(session);
                    }
                } catch (JSchException | IOException e) {
                    log.error("webssh连接异常");
                    log.error("异常信息:{}", e.getMessage());
                    try {
                        //发送错误信息
                        sendMessage(session, ("ERROR : "+e.getMessage()).getBytes());
                    } catch (IOException ex) {
                        log.error("消息发送失败");
                        log.error("异常信息:{}", ex.getMessage());
                    }
                    close(session);
                }
            });
        } else if (ConstantPool.WEBSSH_OPERATE_COMMAND.equals(webSSHData.getOperate())) {
            String command = webSSHData.getCommand();
            SSHConnectInfo sshConnectInfo = clients.get(id);
            if (sshConnectInfo != null) {
                try {
                    ChannelShell channel = (ChannelShell) sshConnectInfo.getChannel();
                    if (channel != null) {
                        channel.setPtySize(webSSHData.getCols(),webSSHData.getRows(),webSSHData.getWidth(),webSSHData.getHeight());
                        transToSSH(channel, command);
                        if (channel.isClosed()) {
                            close(session);
                        }
                    }
                } catch (IOException e) {
                    log.error("webssh连接异常");
                    log.error("异常信息:{}", e.getMessage());
                    try {
                        //发送错误信息
                        sendMessage(session, ("ERROR : "+e.getMessage()).getBytes());
                    } catch (IOException ex) {
                        log.error("消息发送失败");
                        log.error("异常信息:{}", ex.getMessage());
                    }
                    close(session);
                }
            }
        } else if (ConstantPool.WEBSSH_OPERATE_HEARTBEAT.equals(webSSHData.getOperate())) {
            //检查心跳
            SSHConnectInfo sshConnectInfo = clients.get(id);
            if (sshConnectInfo != null) {
                try {
                    //处于连接状态则发送健康数据，不能为空，空则断开连接。
                    if (sshConnectInfo.getChannel().isConnected())
                        sendMessage(session, "Heartbeat healthy".getBytes());
                } catch (IOException e) {
                    log.error("消息发送失败");
                    log.error("异常信息:{}", e.getMessage());
                }
            }
        } else {
            log.error("不支持的操作");
            close(session);
        }
    }

    @OnClose
    public void onClose(Session session) {
       clients.remove(session.getId());
    }


    /**
     * @Description: 使用jsch连接终端
     * @Param: [cloudSSH, webSSHData, webSocketSession]
     * @return: void
     * @Author: NoCortY
     * @Date: 2020/3/7
     */
    private void connectToSSH(SSHConnectInfo sshConnectInfo, WebSSHData webSSHData, Session webSocketSession) throws JSchException, IOException {
        com.jcraft.jsch.Session session = null;
        Properties config = new Properties();
        config.put("StrictHostKeyChecking", "no");
        //获取jsch的会话
        session = sshConnectInfo.getJSch().getSession(webSSHData.getUsername(), webSSHData.getHost(), webSSHData.getPort());
        session.setConfig(config);
        //设置密码
        session.setPassword(webSSHData.getPassword());
        //连接  超时时间30s
        session.connect(30000);

        //开启shell通道
        Channel channels = session.openChannel("shell");
        ChannelShell channel = (ChannelShell) channels;
        channel.setPtySize(webSSHData.getCols(),webSSHData.getRows(),webSSHData.getWidth(),webSSHData.getHeight());

        //通道连接 超时时间3s
        channel.connect(3000);

        //设置channel
        sshConnectInfo.setChannel(channel);

        //读取终端返回的信息流
        try (InputStream inputStream = channel.getInputStream()) {
            //循环读取
            byte[] buffer = new byte[1024];
            int i = 0;
            //如果没有数据来，线程会一直阻塞在这个地方等待数据。
            while ((i = inputStream.read(buffer)) != -1) {
                log.info("返回数据： {}", new String(buffer, 0, i));
                sendMessage(webSocketSession, Arrays.copyOfRange(buffer, 0, i));
            }

        } finally {
            //断开连接后关闭会话
            session.disconnect();
            channel.disconnect();
        }

    }

    /**
     * @Description: 将消息转发到终端
     * @Param: [channel, data]
     * @return: void
     * @Author: NoCortY
     * @Date: 2020/3/7
     */
    private void transToSSH(Channel channel, String command) throws IOException {
        if (channel != null) {
            OutputStream outputStream = channel.getOutputStream();
            outputStream.write(command.getBytes());
            outputStream.flush();
        }
    }

    public void sendMessage(Session session, byte[] buffer) throws IOException {
        session.getBasicRemote().sendText(new String(buffer));
    }

    public void close(Session session) {
        SSHConnectInfo sshConnectInfo = clients.get(session.getId());
        if (sshConnectInfo != null) {
            //断开连接
            if (sshConnectInfo.getChannel() != null) sshConnectInfo.getChannel().disconnect();
            //map中移除
            clients.remove(session.getId());
        }
        try {
            session.close();
        } catch (IOException e) {
            log.error("发生错误", e);
        }
    }
}
